import sys
import time
import cv2
import gymnasium as gym
from tqdm import tqdm
import numpy as np
import copy
import pickle as pkl
import datetime
import os
import keyboard
import rtde_receive

from ft300 import RobotiqFt300
from gym_pih.hardware.Imagingsource_for_windows import Camera as ImagingCamera
from gym_pih.hardware.RealsenseCamera import Camera as RealsenseCamera
from gym_pih.hardware.calculate_slope import calculate_slope
from gym_pih.hardware.LightSource import LightSource
from gym_pih.reward_classifier.load_classifier_for_v3 import Classifier

ip="192.168.3.140"
ft=RobotiqFt300()
realsense_cam = RealsenseCamera()
imaging_cam = ImagingCamera()
realsense_cam.Open()
imaging_cam.Open()
imaging_cam.SetProperty("Gain","Value",0)
imaging_cam.SetProperty("Exposure","Value",float(1/30))
# imaging_cam.SetProperty("Gain","Auto",0)
# imaging_cam.SetProperty("Exposure","Auto",0)
rtde_r=rtde_receive.RTDEReceiveInterface(ip)
light = LightSource()
light.Light_Open(6)
classifier = Classifier()

def run():

    transitions = []
    success_count = 0
    success_needed = 100
    total_count = 0
    pbar = tqdm(total=success_needed)
    
    target_pos=np.array([0.6872853048650824, -0.1308957273597933, 0.17394818223574388, 6.256728706533603e-05, -3.141092053034699, 0.0001309061484982937])
    boundary=np.zeros((2,3))#用于限制机器人的运动范围
    boundary[0]=target_pos[:3]-0.02
    boundary[0][2]=target_pos[2]
    boundary[1]=target_pos[:3]+0.02

    uuid = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    # file_name = f"peg_insert_fail_{success_needed}_demos_{uuid}.pkl"
    file_name = f"peg_insert_fail_{success_needed}_{uuid}.pkl"
    # file_dir = os.path.dirname(os.path.realpath(__file__))  # same dir as this script
    file_dir = r"E:\study\serl_classifier\5"
    file_path = os.path.join(file_dir, file_name)

    

    if not os.path.exists(file_dir):
        os.mkdir(file_dir)
    if os.path.exists(file_path):
        raise FileExistsError(f"{file_name} already exists in {file_dir}")
    if not os.access(file_dir, os.W_OK):
        raise PermissionError(f"No permission to write to {file_dir}")
    
    cv2.namedWindow("Wrist Image", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Wrist Image", 320, 320)
    cv2.namedWindow("Side Image", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Side Image", 320, 320)
    cv2.namedWindow("Slope Image", cv2.WINDOW_NORMAL)



    while success_count < success_needed:
        _,wrist_image=realsense_cam.GetImage()
        _,side_image=imaging_cam.GetImage()
        slope_image=side_image[130:,250:1020]
        _,slope,binary_image = calculate_slope(slope_image,True)
        wrist_image=wrist_image[150:214,215:279]
        side_image=side_image[60:700, 260:900]
        side_image=cv2.resize(side_image,(64,64))

        obs_for_classifier={"images":{ "wrist":wrist_image,"side":side_image}}
        reward=classifier.predict(obs_for_classifier)
        sys.stdout.write(f"\r{reward}reward; {slope}slope;")
        sys.stdout.flush()

        if keyboard.is_pressed("s"):
            pos=np.array(rtde_r.getActualTCPPose()[:3])
            pos=np.clip(pos,boundary[0],boundary[1])
            pos=(pos-boundary[0])/(boundary[1]-boundary[0])
            pos=pos.astype(np.float32)
            force=np.clip(ft.GetForce()/10,-1,1)[:3]
            force=force.astype(np.float32)
            slope=np.array([slope],dtype=np.float32)/90
            obs={"state":np.hstack((pos,force,slope)),
                 "wrist":wrist_image.astype(np.uint8),
                 "side":side_image.astype(np.uint8),}
            next_obs = copy.deepcopy(obs)
            actions=np.zeros(3)
            rew=1
            done=False


            transition = copy.deepcopy(
                dict(
                    observations=obs,
                    actions=actions,
                    next_observations=next_obs,
                    rewards=rew,
                    masks=1.0 - done,
                    dones=done,
                )
            )
        
            transitions.append(transition)
            success_count += rew
            total_count += 1
            print(
                f"{rew}\tGot {success_count} successes of {total_count} trials. {success_needed} successes needed."
            )
            pbar.update(rew)
            time.sleep(0.5)
        if keyboard.is_pressed("g"):
            #concatenate the wrist image and side image, save the image for paper
            concatenated_image=np.concatenate((wrist_image,side_image),axis=0)
            cv2.imwrite(f"success_{total_count}.jpg",concatenated_image)
            total_count+=1
            print("image has been saved")
            time.sleep(0.5)
        cv2.imshow("Wrist Image", wrist_image)
        cv2.imshow("Side Image", side_image)
        cv2.imshow("Slope Image", binary_image)
        cv2.waitKey(1)

    with open(file_path, "wb") as f:
        pkl.dump(transitions, f)
        print(
            f"saved {success_needed} demos and {len(transitions)} transitions to {file_path}"
        )


if __name__ == "__main__":
    try:
        run()
    except KeyboardInterrupt:
        print("KeyboardInterrupt")
    except Exception as e:
        # 捕获所有异常，并打印错误信息
        print(f"发生错误：{e}")
    finally:
        realsense_cam.Close()
        imaging_cam.Close()
        rtde_r.disconnect()
        ft.DisConnect()
        cv2.destroyAllWindows()
        light.Light_Close()
        sys.exit(0)

